#include "core/main.h"

void ValueIterationLearner::learnValueFn(const ReplayCollection &replays, const GameModel &model, ValueFunction &valueFn)
{
    valueFn._transform = _transform;

    // For convenience copy pointers to all frames into a vector and precompute feature transforms for all frames
    vector<const Frame*> allFrames;
    vector<const vector<double> > allFeatureVectors;
    int featureDim = 0; // Dimension of the feature transform of the states
    cout << "Precomputing feature vectors for all frames ... " << endl;
    int frameCount = 0;
    for (auto rit = replays.Replays().begin(); rit != replays.Replays().end(); ++rit) {
        const Replay &replay = **rit;
        for (auto fit = replay.Frames().begin(); fit != replay.Frames().end(); ++fit) {
            if (frameCount % 100 == 0) cout << frameCount << endl;
            frameCount++;

            const Frame &frame = *fit;
            allFrames.push_back(&frame);

            vector<double> featureVector;
            _transform.transform(frame, featureVector);
            featureDim = featureVector.size();
            allFeatureVectors.push_back(featureVector);
        }
    }

    // vector<ControllerState> allActions = ControllerState::getAllStates();
    vector<ControllerState> allActions = ControllerState::getDirectionStates();

    // Also precompute all next states for all actions
    cout << "Computing all next frames ... " << endl;
    vector<vector<Frame> > allNextFrames(allFrames.size());
    frameCount = 0;
    for (UINT i = 0; i < allFrames.size(); i++) {
        if (frameCount % 100 == 0) cout << "Starting frame " << frameCount << endl;
        frameCount++;

        // Make a copy of the current frame
        Frame currentFrame(*allFrames[i]);

        // Stuff it into a history vector
        vector<const Frame*> history;
        history.push_back(&currentFrame);

        for (UINT j = 0; j < allActions.size(); j++) {
            const ControllerState &action = allActions[j];
            currentFrame.controller = action;

            // Predict the next frame
            Frame nextFrame;
            model.AdvanceFrame(history, nextFrame);

            allNextFrames[i].push_back(nextFrame);
        }
    }

    // Initialize the weights of the value function to zero.
    valueFn._weights.clear();
    valueFn._weights.insert(valueFn._weights.begin(), featureDim, 0);

   for (int iterationNum = 0; iterationNum < _params.maxIterations; iterationNum++) {
        cout << "Starting value iteration number " << iterationNum << endl;

        vector<double> ys;
        frameCount = 0;
        for (UINT i = 0; i < allFrames.size(); i++) {
            if (frameCount % 200 == 0) cout << "Starting frame " << frameCount << endl;
            frameCount++;

            // Make a copy of the current frame
            // Frame currentFrame(*allFrames[i]);
            const Frame &currentFrame = *allFrames[i];

            // Stuff it into a history vector
            // vector<const Frame*> history;
            // history.push_back(&currentFrame);

            // Our estimate of the true V(s) for the current frame s
            double y = 0;

            for (UINT j = 0; j < allActions.size(); j++) {
                // const ControllerState &action = allActions[j];
                // currentFrame.controller = action;

                // Predict the next frame
                // Frame nextFrame;
                // model.AdvanceFrame(history, nextFrame);

                // Compute the estimated value of taking this action
                const Frame &nextFrame = allNextFrames[i][j];
                double curReward = rewardFunction(currentFrame);
                double nextValue = valueFn.getValue(nextFrame);
                double val = curReward + _params.discountFactor * nextValue;

                if (val > y) y = val;
            }

            ys.push_back(y);
        }

        // Dammit ... we need linear regression here!!
        // vector<double> newWeights = LinearRegression::linearRegressionBatch(allFeatureVectors, ys, 0.05);
        // vector<double> newWeights = linearRegressionBatch(allFeatureVectors, ys, 0.00001);
        cout << "Starting linear regression ..." << endl;
#ifdef USE_EIGEN
        vector<double> newWeights = linearRegressionEigen(allFeatureVectors, ys);
#else
        vector<double> newWeights = linearRegressionMatrix(allFeatureVectors, ys);
#endif

        // Compute the distance between the old weight vector and the new
        double dist = 0;
        for (UINT i = 0; i < newWeights.size(); i++) {
            double d = newWeights[i] - valueFn._weights[i];
            dist += d * d;
        }
        dist = sqrt(dist);
        cout << "Change in weight vector is " << dist << endl;

        valueFn._weights = newWeights;
    }
}

double ValueIterationLearner::rewardFunction(const Frame &frame)
{
    int index = 2;
    return frame.observedInts[index];
}